module obs_amdar_mod

  use datetime
  use container
  use flogger
  use fiona
  use string
  use const_mod
  use data_translators_utils_mod
  use amdar_mod
  use amdar_cimiss_xml_mod
  use const_mod
  use namelist_mod
  use latlon_process_mod
  use latlon_parallel_mod
  use obs_operator_types_mod
  use obs_err_mod
  use efso_utils_mod

  implicit none

  private

  public obs_amdar_type
  public amdar_record_type

  type obs_amdar_type
    integer :: obs_size = 0
    type(hash_table_type), allocatable :: platforms
    type(linked_list_type), allocatable :: records                    ! Yo
    type(obs_grid_connector_type), allocatable :: connectors(:)
    type(amdar_record_type), allocatable :: simulated_mean_records(:) ! H mean(X)
    type(amdar_record_type), allocatable :: mean_simulated_records(:) ! mean(Y)
    type(amdar_record_type), allocatable :: simulated_records(:,:)    ! Y or Y'
    type(amdar_record_type), allocatable :: innov_records(:)          ! Yo - H mean(X)
    type(obs_err_type), allocatable :: obs_err(:)
    integer, allocatable :: mask(:)
  contains
    procedure :: init => obs_amdar_init
    procedure :: read => obs_amdar_read
    procedure :: filter_by_latlon_and_time => obs_amdar_filter_by_latlon_and_time
    procedure :: clear => obs_amdar_clear
    procedure :: calc_mean_simulated_records => obs_amdar_calc_mean_simulated_records
    procedure :: calc_perb_simulated_records => obs_amdar_calc_perb_simulated_records
    procedure :: calc_innov => obs_amdar_calc_innov
    procedure :: init_write => obs_amdar_init_write
    procedure :: write => obs_amdar_write
    final :: obs_amdar_final
  end type obs_amdar_type

contains

  subroutine obs_amdar_init(this)

    class(obs_amdar_type), intent(inout) :: this

    call this%clear()

    allocate(this%platforms)
    allocate(this%records  )

    this%platforms = hash_table(chunk_size=100000)

  end subroutine obs_amdar_init

  subroutine obs_amdar_read(this, file_path, mute)

    class(obs_amdar_type), intent(inout) :: this
    character(*), intent(in) :: file_path
    logical, intent(in) :: mute

    call amdar_cimiss_xml_read(file_path, this%platforms, this%records, mute)

  end subroutine obs_amdar_read

  subroutine obs_amdar_filter_by_latlon_and_time(this, start_lon, end_lon, start_lat, end_lat, start_time, end_time)

    class(obs_amdar_type), intent(inout) :: this
    real(r8), intent(in) :: start_lon
    real(r8), intent(in) :: end_lon
    real(r8), intent(in) :: start_lat
    real(r8), intent(in) :: end_lat
    type(datetime_type), intent(in) :: start_time
    type(datetime_type), intent(in) :: end_time

    type(linked_list_iterator_type) it
    integer i, k

    it = linked_list_iterator(this%records)
    do while (.not. it%ended())
      select type (record => it%value)
      type is (amdar_record_type)
        if (record%lon < start_lon * deg .or. record%lon > end_lon * deg .or. &
            record%lat < start_lat * deg .or. record%lat > end_lat * deg .or. &
            record%time < start_time .or. record%time > end_time) then
          call this%records%remove_item(it%item)
        end if
      end select
      call it%next()
    end do

    if (parallel_ensemble .and. proc%is_root_proc()) then
      call log_notice('After filter there is ' // to_str(this%records%size) // ' amdar records.')
    end if

    if (allocated(this%connectors)) deallocate(this%connectors)
    allocate(this%connectors(this%records%size))
    if (allocated(this%simulated_mean_records)) deallocate(this%simulated_mean_records)
    allocate(this%simulated_mean_records(this%records%size))
    if (allocated(this%mean_simulated_records)) deallocate(this%mean_simulated_records)
    allocate(this%mean_simulated_records(this%records%size))
    if (allocated(this%simulated_records)) deallocate(this%simulated_records)
    allocate(this%simulated_records(this%records%size,ens_ibeg:ens_iend))
    if (allocated(this%innov_records)) deallocate(this%innov_records)
    allocate(this%innov_records(this%records%size))

    this%obs_size = this%records%size

    ! Initialize observation error structure.
    allocate(this%obs_err(this%obs_size))
    it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => it%value)
      type is (amdar_record_type)
        do k = 1, size(amdar_err_plev)
          if (is_missing(record%p)) exit
          if (amdar_err_plev_bnds(1,k) <= record%p .and. record%p < amdar_err_plev_bnds(2,k)) exit
        end do
        if (k > size(amdar_err_plev)) then
          record%ua = real_missing_value
          record%va = real_missing_value
          record%ta = real_missing_value
          record%sh = real_missing_value
        else
          this%obs_err(i)%u = u_err_amdar(k)
          this%obs_err(i)%v = v_err_amdar(k)
          this%obs_err(i)%t = t_err_amdar(k)
          this%obs_err(i)%q = q_err_amdar(k)
          this%obs_err(i)%p = p_err_amdar(k)
        end if
      end select
      call it%next()
    end do

    ! Set mask for missing values.
    allocate(this%mask(this%obs_size * 4))
    it = linked_list_iterator(this%records)
    i = 1
    do while (.not. it%ended())
      select type (record => it%value)
      type is (amdar_record_type)
        this%mask(i) = merge(0, 1, is_missing(record%ua)); i = i + 1
        this%mask(i) = merge(0, 1, is_missing(record%va)); i = i + 1
        this%mask(i) = merge(0, 1, is_missing(record%ta)); i = i + 1
        this%mask(i) = merge(0, 1, is_missing(record%sh)); i = i + 1
      end select
      call it%next()
    end do

  end subroutine obs_amdar_filter_by_latlon_and_time

  subroutine obs_amdar_clear(this)

    class(obs_amdar_type), intent(inout) :: this

    if (allocated(this%platforms        )) deallocate(this%platforms        )
    if (allocated(this%records          )) deallocate(this%records          )
    if (allocated(this%connectors       )) deallocate(this%connectors       )
    if (allocated(this%simulated_mean_records)) deallocate(this%simulated_mean_records)
    if (allocated(this%mean_simulated_records)) deallocate(this%mean_simulated_records)
    if (allocated(this%simulated_records)) deallocate(this%simulated_records)
    if (allocated(this%innov_records    )) deallocate(this%innov_records    )
    if (allocated(this%obs_err          )) deallocate(this%obs_err          )
    if (allocated(this%mask             )) deallocate(this%mask             )

  end subroutine obs_amdar_clear

  subroutine obs_amdar_calc_mean_simulated_records(this)

    class(obs_amdar_type), intent(inout) :: this

    real(r8) tmp(this%records%size)
    integer i, iens

    do i = 1, this%records%size
      this%mean_simulated_records(i)%ta = 0
      this%mean_simulated_records(i)%p  = 0
      this%mean_simulated_records(i)%ua = 0
      this%mean_simulated_records(i)%va = 0
      this%mean_simulated_records(i)%sh = 0
    end do
    do iens = ens_ibeg, ens_iend
      do i = 1, this%records%size
        this%mean_simulated_records(i)%ta = add(this%simulated_records(i,iens)%ta, this%mean_simulated_records(i)%ta)
        this%mean_simulated_records(i)%p  = add(this%simulated_records(i,iens)%p , this%mean_simulated_records(i)%p )
        this%mean_simulated_records(i)%ua = add(this%simulated_records(i,iens)%ua, this%mean_simulated_records(i)%ua)
        this%mean_simulated_records(i)%va = add(this%simulated_records(i,iens)%va, this%mean_simulated_records(i)%va)
        this%mean_simulated_records(i)%sh = add(this%simulated_records(i,iens)%sh, this%mean_simulated_records(i)%sh)
      end do
    end do

    ! Communicate with other processes.
    do i = 1, this%records%size
      tmp(i) = this%mean_simulated_records(i)%ta
    end do
    call global_sum(proc%comm, tmp)
    do i = 1, this%records%size
      this%mean_simulated_records(i)%ta = tmp(i) / ensemble_size
    end do
    do i = 1, this%records%size
      tmp(i) = this%mean_simulated_records(i)%p
    end do
    call global_sum(proc%comm, tmp)
    do i = 1, this%records%size
      this%mean_simulated_records(i)%p  = tmp(i) / ensemble_size
    end do
    do i = 1, this%records%size
      tmp(i) = this%mean_simulated_records(i)%ua
    end do
    call global_sum(proc%comm, tmp)
    do i = 1, this%records%size
      this%mean_simulated_records(i)%ua = tmp(i) / ensemble_size
    end do
    do i = 1, this%records%size
      tmp(i) = this%mean_simulated_records(i)%va
    end do
    call global_sum(proc%comm, tmp)
    do i = 1, this%records%size
      this%mean_simulated_records(i)%va = tmp(i) / ensemble_size
    end do
    do i = 1, this%records%size
      tmp(i) = this%mean_simulated_records(i)%sh
    end do
    call global_sum(proc%comm, tmp)
    do i = 1, this%records%size
      this%mean_simulated_records(i)%sh = tmp(i) / ensemble_size
    end do

  end subroutine obs_amdar_calc_mean_simulated_records

  subroutine obs_amdar_calc_perb_simulated_records(this)

    class(obs_amdar_type), intent(inout) :: this

    integer i, iens

    ! Perturbation = Y - mean(Y)
    do iens = ens_ibeg, ens_iend
      do i = 1, this%records%size
        this%simulated_records(i,iens)%ta = subtract(this%simulated_records(i,iens)%ta, this%mean_simulated_records(i)%ta)
        this%simulated_records(i,iens)%p  = subtract(this%simulated_records(i,iens)%p , this%mean_simulated_records(i)%p )
        this%simulated_records(i,iens)%ua = subtract(this%simulated_records(i,iens)%ua, this%mean_simulated_records(i)%ua)
        this%simulated_records(i,iens)%va = subtract(this%simulated_records(i,iens)%va, this%mean_simulated_records(i)%va)
        this%simulated_records(i,iens)%sh = subtract(this%simulated_records(i,iens)%sh, this%mean_simulated_records(i)%sh)
      end do
    end do

  end subroutine obs_amdar_calc_perb_simulated_records

  subroutine obs_amdar_calc_innov(this)

    class(obs_amdar_type), intent(inout) :: this

    type(linked_list_iterator_type) it
    integer i, j

    ! Innovation = Yo - H mean(X)
    it = linked_list_iterator(this%records)
    j = 0
    do i = 1, this%records%size
      select type (record => it%value)
      type is (amdar_record_type)
        this%innov_records(i)%ta = subtract(record%ta, this%simulated_mean_records(i)%ta)
        this%innov_records(i)%p  = subtract(record%p , this%simulated_mean_records(i)%p )
        this%innov_records(i)%ua = subtract(record%ua, this%simulated_mean_records(i)%ua)
        this%innov_records(i)%va = subtract(record%va, this%simulated_mean_records(i)%va)
        this%innov_records(i)%sh = subtract(record%sh, this%simulated_mean_records(i)%sh)
        j = j + 1; call check_omb_ua(this%innov_records(i)%ua, this%mask(j))
        j = j + 1; call check_omb_va(this%innov_records(i)%va, this%mask(j))
        j = j + 1; call check_omb_ta(this%innov_records(i)%ta, this%mask(j))
        j = j + 1; call check_omb_sh(this%innov_records(i)%sh, this%mask(j))
      end select
      call it%next()
    end do

  end subroutine obs_amdar_calc_innov

  subroutine obs_amdar_init_write(this, key)

    class(obs_amdar_type), intent(in) :: this
    character(*), intent(in) :: key

    call fiona_add_dim(key, 'amdar_flight'    , size=this%platforms%size)
    call fiona_add_dim(key, 'amdar_record'    , size=this%records%size)
    call fiona_add_var(key, 'amdar_sid'       , long_name='amdar flight id'        , units=''                      , dim_names=['sid_len     ','amdar_flight'], data_type='c')
    call fiona_add_var(key, 'amdar_lon'       , long_name='amdar flight longitude' , units='degree_east'           , dim_names=['amdar_record'])
    call fiona_add_var(key, 'amdar_lat'       , long_name='amdar flight latitude'  , units='degree_north'          , dim_names=['amdar_record'])
    call fiona_add_var(key, 'amdar_alt'       , long_name='amdar flight altitude'  , units='m'                     , dim_names=['amdar_record'])
    call fiona_add_var(key, 'amdar_record_sid', long_name='amdar record flight id' , units=''                      , dim_names=['sid_len     ','amdar_record'], data_type='c')
    call fiona_add_var(key, 'amdar_time'      , long_name='amdar observation time' , units='hours since 1970-01-01', dim_names=['amdar_record'], data_type='r8')
    call fiona_add_var(key, 'amdar_ta'        , long_name='amdar air temperature'  , units='degC'                  , dim_names=['amdar_record'], data_type='r4', missing_value=real_missing_value)
    call fiona_add_var(key, 'amdar_ua'        , long_name='amdar u wind component' , units='m s-1'                 , dim_names=['amdar_record'], data_type='r4', missing_value=real_missing_value)
    call fiona_add_var(key, 'amdar_va'        , long_name='amdar v wind component' , units='m s-1'                 , dim_names=['amdar_record'], data_type='r4', missing_value=real_missing_value)
    call fiona_add_var(key, 'amdar_sh'        , long_name='amdar specific humidity', units='g g-1'                 , dim_names=['amdar_record'], data_type='r4', missing_value=real_missing_value)

    ! Simulated observation
    call fiona_add_var(key, 'amdar_ta_sim'    , long_name='simulated amdar_ta'     , units='degC'                  , dim_names=['amdar_record','ens         '], data_type='r4', missing_value=real_missing_value)
    call fiona_add_var(key, 'amdar_ua_sim'    , long_name='simulated amdar_ua'     , units='m s-1'                 , dim_names=['amdar_record','ens         '], data_type='r4', missing_value=real_missing_value)
    call fiona_add_var(key, 'amdar_va_sim'    , long_name='simulated amdar_va'     , units='m s-1'                 , dim_names=['amdar_record','ens         '], data_type='r4', missing_value=real_missing_value)
    call fiona_add_var(key, 'amdar_sh_sim'    , long_name='simulated amdar_sh'     , units='g g-1'                 , dim_names=['amdar_record','ens         '], data_type='r4', missing_value=real_missing_value)

  end subroutine obs_amdar_init_write

  subroutine obs_amdar_write(this, key)

    class(obs_amdar_type), intent(in) :: this
    character(*), intent(in) :: key

    type(hash_table_iterator_type) flight_it
    type(linked_list_iterator_type) record_it
    character(8), allocatable :: cvar(:)
    real(r8), allocatable :: rvar(:), rvar_ens(:,:)
    integer i, iens

    allocate(cvar(this%platforms%size))

    flight_it = hash_table_iterator(this%platforms)
    do i = 1, this%platforms%size
      select type (flight => flight_it%value)
      type is (amdar_flight_type)
        cvar(i) = flight%name
      end select
      call flight_it%next()
    end do
    call fiona_output(key, 'amdar_sid', cvar)

    deallocate(cvar)

    allocate(rvar(this%records%size))

    record_it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => record_it%value)
      type is (amdar_record_type)
        rvar(i) = record%lon
      end select
      call record_it%next()
    end do
    call fiona_output(key, 'amdar_lon', rvar)

    record_it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => record_it%value)
      type is (amdar_record_type)
        rvar(i) = record%lat
      end select
      call record_it%next()
    end do
    call fiona_output(key, 'amdar_lat', rvar)

    record_it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => record_it%value)
      type is (amdar_record_type)
        rvar(i) = record%h
      end select
      call record_it%next()
    end do
    call fiona_output(key, 'amdar_alt', rvar)

    record_it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => record_it%value)
      type is (amdar_record_type)
        rvar(i) = record%time%timestamp() / 3600.0d0
      end select
      call record_it%next()
    end do
    call fiona_output(key, 'amdar_time', rvar)

    record_it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => record_it%value)
      type is (amdar_record_type)
        rvar(i) = record%ta
      end select
      call record_it%next()
    end do
    call fiona_output(key, 'amdar_ta', rvar)

    record_it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => record_it%value)
      type is (amdar_record_type)
        rvar(i) = record%ua
      end select
      call record_it%next()
    end do
    call fiona_output(key, 'amdar_ua', rvar)

    record_it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => record_it%value)
      type is (amdar_record_type)
        rvar(i) = record%va
      end select
      call record_it%next()
    end do
    call fiona_output(key, 'amdar_va', rvar)

    record_it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => record_it%value)
      type is (amdar_record_type)
        rvar(i) = record%sh
      end select
      call record_it%next()
    end do
    call fiona_output(key, 'amdar_sh', rvar)

    deallocate(rvar)

    allocate(cvar(this%records%size))

    record_it = linked_list_iterator(this%records)
    do i = 1, this%records%size
      select type (record => record_it%value)
      type is (amdar_record_type)
        cvar(i) = record%flight%name
      end select
      call record_it%next()
    end do
    call fiona_output(key, 'amdar_record_sid', cvar)

    deallocate(cvar)

    allocate(rvar_ens(size(this%simulated_records, 1), &
                      lbound(this%simulated_records, 2): &
                      ubound(this%simulated_records, 2)))

    do iens = lbound(this%simulated_records, 2), ubound(this%simulated_records, 2)
      do i = 1, size(this%simulated_records, 1)
        rvar_ens(i,iens) = this%simulated_records(i,iens)%ta
      end do
    end do
    call fiona_output(key, 'amdar_ta_sim', rvar_ens)

    do iens = lbound(this%simulated_records, 2), ubound(this%simulated_records, 2)
      do i = 1, size(this%simulated_records, 1)
        rvar_ens(i,iens) = this%simulated_records(i,iens)%ua
      end do
    end do
    call fiona_output(key, 'amdar_ua_sim', rvar_ens)

    do iens = lbound(this%simulated_records, 2), ubound(this%simulated_records, 2)
      do i = 1, size(this%simulated_records, 1)
        rvar_ens(i,iens) = this%simulated_records(i,iens)%va
      end do
    end do
    call fiona_output(key, 'amdar_va_sim', rvar_ens)

    do iens = lbound(this%simulated_records, 2), ubound(this%simulated_records, 2)
      do i = 1, size(this%simulated_records, 1)
        rvar_ens(i,iens) = this%simulated_records(i,iens)%va
      end do
    end do
    call fiona_output(key, 'amdar_va_sim', rvar_ens)

    deallocate(rvar_ens)

  end subroutine obs_amdar_write

  subroutine obs_amdar_final(this)

    type(obs_amdar_type), intent(inout) :: this

    call this%clear()

  end subroutine obs_amdar_final

end module obs_amdar_mod
